'''
 * @Author: codepanda
 * @Date: 2023-12-21 16:11:21
LastEditors: CodePanda
 * @Email: codepanda@qq.com
LastEditTime: 2023-12-22 21:28:53
'''

import numpy as np
import pandas as pd
from pathlib import Path
from obspy import read

from obspy import read, Stream
from obspy.io.sac import SACTrace
from obspy.taup import TauPyModel

import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
from matplotlib import colors, image

from configparser import ConfigParser


def hk(station_name:str, st: Stream, config: ConfigParser, *args, **kwargs):
    """_summary_

    Args:
        station_name (str): station name
        st (Stream): P receiver function from SAC file to obspy.Stream
        config (ConfigParser): config information from config.ini
    Raises:
        Exception: raise Exception if trace is not has attribute 'sac' 

    Returns:
        _list_: Ps H-k image, PpPs H-k image, PsPs+PpSs H-k image, The weighted average of the above values H-k image, best point of H and k, contour line for level 1 of H-k image, best point of H and k residual.
    """
    model_name = config.get("hk", "model_name")
    h_min = config.getfloat("hk", "hmin")
    h_max = config.getfloat("hk", "hmax")
    k_min = config.getfloat("hk", "kmin")
    k_max = config.getfloat("hk", "kmax")
    num = config.getint("hk", "num")
    vp = config.getfloat("hk", "vp")
    pre_time = config.getfloat("hk", "pre_time")
    post_time = config.getfloat("hk", "post_time")
    w1 = config.getfloat("hk", "w1")
    w2 = config.getfloat("hk", "w2")
    w3 = config.getfloat("hk", "w3")
    
    rps_list = []
    rppps_list = []
    rppss_list = []
    
    n = len(st)
    h = np.linspace(h_min, h_max, num=num)
    k = np.linspace(k_min, k_max, num=num)
    h_grids, k_grids = np.meshgrid(h, k)
    vs = vp / k
    
    for tr in st:
        if hasattr(tr.stats, "sac"):
            sac_header = tr.stats.sac
        else:
            raise Exception("The read is not a SAC file or does not contain SAC header variables.")
        
        gcarc = sac_header.gcarc
        dist = sac_header.dist
        b = sac_header.b
        evdp = sac_header.evdp
        delta = tr.stats.delta
        # 根据震源深度和震中距计算射线参数，并将单位从（s/°）转换为 （s/km）
        model = TauPyModel(model=model_name)
        ray_info = model.get_ray_paths(source_depth_in_km=evdp, distance_in_degree=gcarc, phase_list=["P",])[0]
        p_ray_parm = ray_info.ray_param_sec_degree / dist * gcarc

        tps = h_grids.T * (np.sqrt(1. / vs**2 - p_ray_parm ** 2) - np.sqrt(1. / vp**2 - p_ray_parm ** 2))
        tppps = h_grids.T * (np.sqrt(1. / vs**2 - p_ray_parm ** 2) + np.sqrt(1. / vp**2 - p_ray_parm ** 2))
        tppss = h_grids.T * np.sqrt(1. / vs**2 - p_ray_parm ** 2) * 2
        
        rps_tmp = np.zeros((num,num))
        rppps_tmp = np.zeros((num,num))
        rppss_tmp = np.zeros((num,num))
        
        for index, t in enumerate(zip(tps, tppps, tppss)):
            tps_index = np.round((t[0]-tr.stats.sac.b)/tr.stats.delta)
            tppps_index = np.round((t[1]-tr.stats.sac.b)/tr.stats.delta)
            tppss_index = np.round((t[2]-tr.stats.sac.b)/tr.stats.delta)

            b_tps = (tps_index >= int(round((pre_time - b) / tr.stats.delta))) & (tps_index <= int(round((post_time - b) / tr.stats.delta)))
            b_tppps = (tppps_index >= int(round((pre_time - b) / tr.stats.delta))) & (tppps_index <= int(round((post_time - b) / tr.stats.delta)))
            b_tppss = (tppss_index >= int(round((pre_time - b) / tr.stats.delta))) & (tppss_index <= int(round((post_time - b) / tr.stats.delta)))
            
            tps_index = tps_index * b_tps
            tppps_index = tppps_index * b_tppps
            tppss_index = tppss_index * b_tppss
            
            for a, i in enumerate(tps_index):
                if int(i) != 0:
                    rps_tmp[index, a] = tr.data[int(i)]
                    
            for a, i in enumerate(tppps_index):
                if int(i) != 0:
                    rppps_tmp[index, a] = tr.data[int(i)]
                    
            for a, i in enumerate(tppss_index):
                if int(i) != 0:
                    rppss_tmp[index, a] = tr.data[int(i)]
            
        rps_list.append(rps_tmp)
        rppps_list.append(rppps_tmp)
        rppss_list.append(rppss_tmp)
        
    rps_narr = np.array(rps_list)
    rppps_narr = np.array(rppps_list)
    rppss_narr = np.array(rppss_list)
    
    rps_avg = (np.sum(rps_narr, axis=0)/n).T
    rppps_avg = (np.sum(rppps_narr, axis=0)/n).T
    rppss_avg = (np.sum(rppss_narr, axis=0)/n).T
        
    S = w1 * rps_avg + w2 * rppps_avg - w3 * rppss_avg
    max_pos= np.where(S == np.max(S))

    best_h = h_grids[max_pos][0]
    best_k = k_grids[max_pos][0]
    # 归一化 rps_avg
    rps_avg = (rps_avg - rps_avg.min()) / (rps_avg.max() - rps_avg.min())
    rppps_avg = (rppps_avg - rppps_avg.min()) / (rppps_avg.max() - rppps_avg.min())
    rppss_avg = (rppss_avg - rppss_avg.min()) / (rppss_avg.max() - rppss_avg.min())
    S = (S - S.min()) / (S.max() - S.min())
    
    cvalue = 1 - np.std(S.reshape(S.size)) / np.sqrt(n)
    cs = plt.contour(h_grids, k_grids, S, [cvalue])
    # cs_path = cs.collections[0].get_paths()[0].vertices
    cs_path = cs.get_paths()[0].vertices
    maxhsig = (np.max(cs_path[:, 0]) - np.min(cs_path[:, 0])) / 2
    maxksig = (np.max(cs_path[:, 1]) - np.min(cs_path[:, 1])) / 2
    plt.close()
    
    out_root = Path(config.get('output', 'out_root'))
    name_dict = {}
    # taup model name
    name_dict['station_name'] = station_name
    name_dict['model_name'] = config.get('hk', 'model_name')
    # Range of H and k
    name_dict['hmin'] = config.get('hk', 'hmin')
    name_dict['hmax'] = config.get('hk', 'hmax')
    name_dict['kmin'] = config.get('hk', 'kmin')
    name_dict['kmax'] = config.get('hk', 'kmax')
    name_dict['num'] = config.get('hk', 'num')
    # Crustal Vp
    name_dict['vp'] = config.get('hk', 'vp')
    
    hK_results_file_name = config.get('output', 'hK_results_file_name').format(**name_dict)
    model_name = config.get('hk', 'model_name')
    
    out_root.mkdir(exist_ok=True, parents=True)
    hK_results_file_path = out_root / hK_results_file_name
    with open(hK_results_file_path, 'w') as f:
        print(f'{station_name}, {best_h:.2f}, {maxhsig:.2f}, {best_k:.2f}, {maxksig:.2f}')
        f.write(f'{station_name}, {best_h:.2f}, {maxhsig:.2f}, {best_k:.2f}, {maxksig:.2f}\n')
    
    return rps_avg, rppps_avg, rppss_avg, S, best_h, best_k, cvalue, maxhsig, maxksig
    
    
def plot_hk(station_name, rps_avg, rppps_avg, rppss_avg, S, best_h, best_k, cvalue, maxhsig, maxksig, config, is_save=True, is_show=False):
    h_min = config.getfloat('hk', 'hmin')
    h_max = config.getfloat('hk', 'hmax')
    k_min = config.getfloat('hk', 'kmin')
    k_max = config.getfloat('hk', 'kmax')
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 8))
    fig.suptitle(f'{station_name}\nMoho depth = ${best_h:.2f}\\pm{maxhsig:.2f}$ km\n$V_P/V_P={best_k:.2f}\\pm{maxksig:.2f}$')
    cmap = "Blues"
    ax1.imshow(rps_avg, aspect='auto', origin='lower', extent=[h_min, h_max, k_min, k_max], cmap=cmap, alpha=0.8)
    ax2.imshow(rppps_avg, aspect='auto', origin='lower', extent=[h_min, h_max, k_min, k_max], cmap=cmap, alpha=0.8)
    ax3.imshow(rppss_avg, aspect='auto', origin='lower', extent=[h_min, h_max, k_min, k_max], cmap=cmap, alpha=0.8)
    im4 = ax4.imshow(S, aspect='auto', origin='lower', extent=[h_min, h_max, k_min, k_max], cmap=cmap, alpha=0.8)
    ax4.contour(S, [cvalue, 1], colors='k', origin='lower', extent=[h_min, h_max, k_min, k_max])

    ax4.plot(best_h, best_k, color='red', marker='s', markerfacecolor='none')
    ax4.axvline(best_h, color='red', linestyle='--', linewidth=0.5)
    ax4.axhline(best_k, color='red', linestyle='--', linewidth=0.5)
    bar = plt.colorbar(im4, ax=[ax1, ax2, ax3, ax4], ticks=[0.0, 0.2, 0.4, 0.6, 0.8, 1.0])

    ax1.set_title("Ps")
    ax2.set_title("PpPs")
    ax3.set_title("PsPs+PpSs")
    ax4.set_title("H-k")

    ax1.set_ylabel("Vp/Vs")
    ax3.set_xlabel("Moho depth(km)")
    ax3.set_ylabel("Vp/Vs")
    ax4.set_xlabel("Moho depth(km)")

    if is_show:
        plt.show()
    
    out_root = Path(config.get('output', 'out_root'))
    name_dict = {}
    # taup model name
    name_dict['station_name'] = station_name
    name_dict['model_name'] = config.get('hk', 'model_name')
    # Range of H and k
    name_dict['hmin'] = config.get('hk', 'hmin')
    name_dict['hmax'] = config.get('hk', 'hmax')
    name_dict['kmin'] = config.get('hk', 'kmin')
    name_dict['kmax'] = config.get('hk', 'kmax')
    name_dict['num'] = config.get('hk', 'num')
    # Crustal Vp
    name_dict['vp'] = config.get('hk', 'vp')
    
    hk_img_name = config.get('output', 'hk_img_name').format(**name_dict)
    model_name = config.get('hk', 'model_name')
    
    out_root.mkdir(exist_ok=True, parents=True)
    hk_img_path = out_root / hk_img_name
    if is_save:
        plt.savefig(str(hk_img_path), dpi=600)
    plt.close()
    

if __name__ == '__main__':
    config = ConfigParser()
    config.read('hk.ini', encoding='UTF-8')
    
    rf_path = Path(config.get('input', 'rf_path'))
    stations_name = config.get('input', 'stations_name')
    best_dir = config.get('input', 'best_dir_name')
    rf_ext = config.get('input', 'rf_ext')
    
    son_dir_list = list(f.name for f in rf_path.iterdir() if f.is_dir())
    if stations_name == '*':
        station_list = son_dir_list
    else:
        station_list = stations_name.split(',')
    
    for station_name in station_list:
        st = Stream()
        if station_name in son_dir_list:
            path_patten=f"./{station_name}/{best_dir}/*.{rf_ext}"
            for f in rf_path.glob(path_patten):
                # print(f)
                st += read(str(f))
                
            if len(st) > 0:
                rps_avg, rppps_avg, rppss_avg, S, h_max, k_max, cvalue, maxhsig, maxksig = hk(station_name=station_name, st=st, config=config)
            
                plot_hk(station_name, rps_avg, rppps_avg, rppss_avg, S, h_max, k_max, cvalue, maxhsig, maxksig, config)
    